#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ChromaDB REST API 服务
提供聊天历史的向量存储和检索功能
"""

import os
import sys

# 必须在导入 chromadb 之前禁用遥测（兼容 Python 3.8）
os.environ['ANONYMIZED_TELEMETRY'] = 'False'
os.environ['CHROMA_TELEMETRY_IMPL'] = 'chromadb.telemetry.product.posthog.Posthog'

# 尝试阻止加载 posthog 模块
class DummyPosthog:
    def __init__(self, *args, **kwargs):
        pass
    def capture(self, *args, **kwargs):
        pass
    def identify(self, *args, **kwargs):
        pass

sys.modules['posthog'] = type('posthog', (), {'Client': DummyPosthog})()

from flask import Flask, request, jsonify
from flask_cors import CORS
import chromadb
from chromadb.config import Settings
import uuid
from datetime import datetime, timedelta, timezone

app = Flask(__name__)
CORS(app)

# 初始化 ChromaDB
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
chroma_path = os.path.join(project_root, 'chroma')

# 创建持久化客户端，禁用遥测以兼容 Python 3.8
client = chromadb.PersistentClient(
    path=chroma_path,
    settings=Settings(
        anonymized_telemetry=False,
        allow_reset=True
    )
)

# 获取或创建集合
collection = client.get_or_create_collection(
    name="chat_history",
    metadata={"description": "风力发电叶片缺陷检测系统聊天历史"}
)

print(f"✅ ChromaDB 已启动")
print(f"📁 数据目录: {chroma_path}")
print(f"📊 当前有 {collection.count()} 条对话记录")

@app.route('/api/health', methods=['GET'])
def health_check():
    """健康检查"""
    return jsonify({
        'status': 'ok',
        'count': collection.count(),
        'path': chroma_path
    })

@app.route('/api/save', methods=['POST'])
def save_conversation():
    """保存对话"""
    try:
        data = request.json
        conversation_id = data.get('id') or str(uuid.uuid4())
        embedding = data.get('embedding')
        metadata = data.get('metadata', {})
        document = data.get('document')
        
        if not embedding or not document:
            return jsonify({'error': '缺少必要参数'}), 400
        
        # 保存到 ChromaDB
        collection.add(
            ids=[conversation_id],
            embeddings=[embedding],
            metadatas=[metadata],
            documents=[document]
        )
        
        print(f"💾 保存对话: {conversation_id[:8]}... (共 {collection.count()} 条)")
        
        return jsonify({
            'success': True,
            'id': conversation_id,
            'count': collection.count()
        })
    except Exception as e:
        print(f"❌ 保存失败: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/search', methods=['POST'])
def search_conversations():
    """语义搜索"""
    try:
        data = request.json
        query_embedding = data.get('embedding')
        limit = data.get('limit', 5)
        
        if not query_embedding:
            return jsonify({'error': '缺少查询向量'}), 400
        
        # 查询
        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=limit
        )
        
        conversations = []
        if results['ids'] and results['ids'][0]:
            for i in range(len(results['ids'][0])):
                conversations.append({
                    'id': results['ids'][0][i],
                    'distance': results['distances'][0][i] if results.get('distances') else None,
                    'metadata': results['metadatas'][0][i] if results.get('metadatas') else {},
                    'document': results['documents'][0][i] if results.get('documents') else ''
                })
        
        return jsonify({'conversations': conversations})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/recent', methods=['GET'])
def get_recent():
    """获取最近对话（按时间排序，支持分页和日期筛选）"""
    try:
        limit = int(request.args.get('limit', 20))
        page = int(request.args.get('page', 1))
        date_filter = request.args.get('dateFilter', 'all')
        
        print(f"📜 [Python API] 获取历史记录 - 页码: {page}, 每页: {limit}, 筛选: {date_filter}")
        
        # 获取所有数据
        total_count = collection.count()
        print(f"📊 [Python API] 数据库中总共有 {total_count} 条记录")
        
        if total_count == 0:
            return jsonify({
                'conversations': [],
                'total': 0,
                'hasMore': False,
                'page': page
            })
        
        # 获取所有记录
        results = collection.get(
            limit=total_count,
            include=['metadatas', 'documents']
        )
        
        conversations = []
        if results['ids']:
            for i in range(len(results['ids'])):
                metadata = results['metadatas'][i] if results.get('metadatas') else {}
                conversations.append({
                    'id': results['ids'][i],
                    'metadata': metadata,
                    'document': results['documents'][i] if results.get('documents') else ''
                })
        
        print(f"📥 [Python API] 从数据库获取到 {len(conversations)} 条记录")
        
        # 日期筛选
        if date_filter != 'all':
            # 获取当前北京时间（UTC+8）
            china_tz = timezone(timedelta(hours=8))
            now_china = datetime.now(china_tz)
            today_china = now_china.date()
            
            print(f"🔍 [Python API] 当前北京时间: {now_china.isoformat()}")
            print(f"🔍 [Python API] 当前日期: {today_china}")
            
            filtered_conversations = []
            
            for conv in conversations:
                timestamp_str = conv['metadata'].get('timestamp', '')
                if not timestamp_str:
                    print(f"⚠️ [Python API] 记录 {conv['id'][:8]} 没有时间戳")
                    continue
                
                try:
                    # 解析时间戳 - 统一转换为北京时间
                    if timestamp_str.endswith('Z'):
                        # UTC 时间，转换为北京时间
                        conv_time_utc = datetime.fromisoformat(timestamp_str.replace('Z', '+00:00'))
                        conv_time_china = conv_time_utc.astimezone(china_tz)
                    elif '+' in timestamp_str or timestamp_str.count('-') > 2:
                        # 已经包含时区信息
                        conv_time_china = datetime.fromisoformat(timestamp_str).astimezone(china_tz)
                    else:
                        # 没有时区信息，假设是 UTC 时间
                        conv_time_utc = datetime.fromisoformat(timestamp_str).replace(tzinfo=timezone.utc)
                        conv_time_china = conv_time_utc.astimezone(china_tz)
                    
                    conv_date_china = conv_time_china.date()
                    
                    print(f"🕐 [Python API] 记录 {conv['id'][:8]} 北京时间: {conv_time_china.isoformat()}, 日期: {conv_date_china}")
                    
                    # 根据筛选条件判断
                    should_include = False
                    
                    if date_filter == 'today':
                        should_include = conv_date_china == today_china
                        print(f"   {'✅' if should_include else '❌'} 今天筛选: {conv_date_china} == {today_china} ? {should_include}")
                    elif date_filter == 'yesterday':
                        yesterday_china = (now_china - timedelta(days=1)).date()
                        should_include = conv_date_china == yesterday_china
                        print(f"   {'✅' if should_include else '❌'} 昨天筛选: {conv_date_china} == {yesterday_china} ? {should_include}")
                    elif date_filter == 'week':
                        week_ago_china = now_china - timedelta(days=7)
                        should_include = conv_time_china >= week_ago_china
                        print(f"   {'✅' if should_include else '❌'} 最近一周筛选")
                    elif date_filter == 'month':
                        month_ago_china = now_china - timedelta(days=30)
                        should_include = conv_time_china >= month_ago_china
                        print(f"   {'✅' if should_include else '❌'} 最近一月筛选")
                    elif date_filter == 'year':
                        should_include = conv_time_china.year == now_china.year
                        print(f"   {'✅' if should_include else '❌'} 今年筛选: {conv_time_china.year} == {now_china.year} ? {should_include}")
                    
                    if should_include:
                        filtered_conversations.append(conv)
                        
                except Exception as e:
                    print(f"⚠️ [Python API] 解析时间戳失败: {timestamp_str}, 错误: {str(e)}")
                    import traceback
                    traceback.print_exc()
                    continue
            
            conversations = filtered_conversations
            print(f"🔍 [Python API] 日期筛选 '{date_filter}' 后剩余 {len(conversations)} 条记录")
        
        # 按时间排序（从新到旧）
        conversations.sort(
            key=lambda x: x['metadata'].get('timestamp', ''), 
            reverse=True
        )
        
        # 分页
        total = len(conversations)
        start_idx = (page - 1) * limit
        end_idx = start_idx + limit
        page_conversations = conversations[start_idx:end_idx]
        has_more = end_idx < total
        
        print(f"✅ [Python API] 返回第 {page} 页的 {len(page_conversations)} 条记录（总共 {total} 条）")
        
        return jsonify({
            'conversations': page_conversations,
            'total': total,
            'hasMore': has_more,
            'page': page
        })
    except Exception as e:
        print(f"❌ [Python API] 获取历史记录失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/api/delete/<conversation_id>', methods=['DELETE'])
def delete_conversation(conversation_id):
    """删除对话"""
    try:
        collection.delete(ids=[conversation_id])
        print(f"🗑️ 删除对话: {conversation_id[:8]}... (剩余 {collection.count()} 条)")
        return jsonify({'success': True, 'count': collection.count()})
    except Exception as e:
        print(f"❌ 删除失败: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/clear', methods=['POST'])
def clear_all():
    """清空所有对话"""
    try:
        global collection
        client.delete_collection(name="chat_history")
        collection = client.create_collection(
            name="chat_history",
            metadata={"description": "风力发电叶片缺陷检测系统聊天历史"}
        )
        print(f"🗑️ 已清空所有对话记录")
        return jsonify({'success': True, 'count': 0})
    except Exception as e:
        print(f"❌ 清空失败: {str(e)}")
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    print(f"\n🚀 ChromaDB API 服务启动在: http://localhost:8001")
    print(f"💡 请在 Node.js 服务中配置: CHROMA_API_URL=http://localhost:8001")
    app.run(host='0.0.0.0', port=8001, debug=False)
